# -*- coding: utf-8 -*-
# created on 2017/8/25

from mathsolver.functions.base import *


class ShuLieA1ValueUpdate001(BaseFunction):
    """
    已知数列{a_{n}}的首项为1,S_{n}为数列{a_{n}}的前n项和,a_{n+1}=a_{n}+1,求{a_{n}}的前n项和S_{n}.
    """
    def solver(self, *args):
        assert len(args) == 2
        assert isinstance(args[0], BaseSequence)
        sl = self.search(args[0].name)
        assert isinstance(args[1], (BaseNumber, BaseVariable))
        value = args[1].sympify()
        sl.aValue = value
        self.output.append(sl)
        return self


class ShuLieA1ValueUpdate002(BaseFunction):
    """
    首项为1,公比为2的等比数列的前4项和S_{4}=.
    """
    def solver(self, *args):
        assert len(args) == 2
        assert isinstance(args[0], BaseSeqPoly)
        p = r'[a-zA-Z]_\(.*?\)'
        m = re.findall(p, args[0].value)
        m = set(m)
        m = list(m)
        assert len(m) == 1
        m = m[0]
        m = sympify(m)
        m_name = m.func
        sl = self.search(str(m_name))
        assert isinstance(args[1], (BaseNumber, BaseVariable))
        value = args[1].sympify()
        sl.aValue = value
        self.output.append(sl)
        return self


class ShuLieA1ValueUpdate003(BaseFunction):
    """
    设正项等比数列{a_{n}}的首项a_{1}=\\frac{1}{2},前n项和为S_{n},2^{10}S_{30}-(2^{10}+1)S_{20}+S_{10}=0,求{a_{n}}的通项.
    """
    def solver(self, *args):
        assert len(args) == 2
        assert isinstance(args[0], BaseSequence)
        sl = self.search(args[0].name)
        assert isinstance(args[1], BaseSeqEq)
        eq = args[1].sympify()
        a_name = sl.aName
        a_name = sympify(a_name)
        f = eq[0] - eq[1]
        value = solve(f, a_name(1))[0]
        sl.aValue = value
        self.output.append(sl)
        return self


class ShuLieA1ValueUpdate(BaseFunction):
    CLS = [ShuLieA1ValueUpdate001, ShuLieA1ValueUpdate002, ShuLieA1ValueUpdate003]

    def solver(self, *args):
        r = None
        for cl in ShuLieA1ValueUpdate.CLS:
            try:
                r = cl(self.known, verbose=True).solver(*args)
                break
            except Exception:
                pass
        if not r:
            raise 'try fail'
        return r
